import torch.nn as nn
import torch.nn.functional as F
from segment_anything import sam_model_registry, SamPredictor

class BaseModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
    
    def setup(self):
        self.model = sam_model_registry[self.cfg['model']['type']](checkpoint=self.cfg['model']['checkpoint'])
        self.model.train()
        if self.cfg['model']['freeze']['image_encoder']:
            for name, param in self.model.image_encoder.named_parameters():
                param.requires_grad = False
            print(f'Freezing image_encoder')

        if self.cfg['model']['freeze']['prompt_encoder']:
            for name, param in self.model.prompt_encoder.named_parameters():
                param.requires_grad = False
            print(f'Freezing prompt_encoder')
            
        # freeze mask decoder 参数
        if self.cfg['model']['freeze']['mask_decoder']:
            for name, param in self.model.mask_decoder.named_parameters():
                param.requires_grad = False
            print(f'Freezing mask_decoder')
    
    def forward(self, images, bboxes):
        _, _, H, W = images.shape
        image_embeddings = self.model.image_encoder(images)
        pred_masks = []
        ious = []
        for embedding, bbox in zip(image_embeddings, bboxes):
            sparse_embeddings, dense_embeddings = self.model.prompt_encoder(
                points=None,
                boxes=bbox,
                masks=None,
            )

            low_res_masks, iou_predictions = self.model.mask_decoder(
                image_embeddings=embedding.unsqueeze(0),
                image_pe=self.model.prompt_encoder.get_dense_pe(),
                sparse_prompt_embeddings=sparse_embeddings,
                dense_prompt_embeddings=dense_embeddings,
                multimask_output=False,
            )

            masks = F.interpolate(
                low_res_masks,
                (H, W),
                mode="bilinear",
                align_corners=False,
            )
            pred_masks.append(masks.squeeze(1))
            ious.append(iou_predictions)

        return pred_masks, ious
    
    def get_predictor(self):
        return SamPredictor(self.model)


class Model(BaseModel):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
    
    def forward(self, images, bboxes, point_coords, point_labels):
        _, _, H, W = images.shape
        image_embeddings = self.model.image_encoder(images)
        pred_masks = []
        ious = []

        for embedding, bbox, coord, label in zip(image_embeddings, bboxes, point_coords, point_labels):
            sparse_embeddings, dense_embeddings = self.model.prompt_encoder(
                points=(coord.unsqueeze(0), label.unsqueeze(0)),
                boxes=bbox.unsqueeze(0),
                masks=None,
            )

            low_res_mask, iou_predictions = self.model.mask_decoder(
                image_embeddings=embedding.unsqueeze(0),
                image_pe = self.model.prompt_encoder.gen_dense_pe(),
                sparse_prompt_embeddings=sparse_embeddings,
                dense_prompt_embeddings=dense_embeddings,
                multimask_output=False,
            )
            masks = F.interpolate(
                low_res_mask,
                size=(H, W),
                mode='bilinear',
                align_corners=False,
            )
            pred_masks.append(masks.sequeeze(0))
            ious.append(iou_predictions)
        return pred_masks, ious

    